# -*- coding: utf-8 -*-
import os
import uuid

from loguru import logger
from pydantic import BaseModel
from fastapi import Request

from schemas.response_entity import JsonModel
from db.repository.conversation import db_add_conversation, db_save_dialog, db_search_dialog
from db.repository.emergency_file import db_get_ep_file_by_id
from db.repository.attributes import db_get_all_accident_types
from utils.file_utils import to_abs_path

from service.emergency_drill.emergency_3 import run_emergency_3
from service.emergency_drill.emergency_32 import run_emergency_32
from service.emergency_drill.emergency_321 import run_emergency_321


class ChatEntity(BaseModel):
    query: str
    token: str
    file_id: uuid.UUID = None
    dialog_id: uuid.UUID = None


async def emergency_chat(request: Request, payload: ChatEntity):
    try:
        user = request.state.user  # 直接获取
        headers = None

        if payload.file_id is None:
            return JsonModel(code=404, data=None, msg=f"请选择文件").to_response(headers=headers)

        file = db_get_ep_file_by_id(payload.file_id)
        if file is None:
            return JsonModel(code=404, data=None, msg=f"文件不存在").to_response(headers=headers)

        file_location = to_abs_path(file.location)
        if not os.path.exists(file_location):
            return JsonModel(code=404, data=None, msg=f"文件不存在").to_response(headers=headers)

        chunks_file_path = file_location.rsplit(".", 1)[0] + "_chunks.json"
        if not os.path.exists(chunks_file_path):
            return JsonModel(code=404, data=None, msg=f"文件不存在").to_response(headers=headers)

        user_id = user.id
        if payload.dialog_id:
            logger.debug(f"user_id: {user_id}")
            dialog = db_search_dialog(dialog_id=payload.dialog_id, user_id=user_id, dialog_type=3)
            if dialog is None:
                return JsonModel(code=404, data=None, msg=f"对话不存在,请更换对话").to_response(headers=headers)
        else:
            dialog = db_save_dialog(payload.query, user.id, dialog_type=3)

        accident_types = db_get_all_accident_types()
        # file.accident_type 可能包含 int 或 str，这里统一转为 int 集合再比对
        raw_ids = file.accident_type or []
        norm_ids = set()
        for v in raw_ids:
            try:
                norm_ids.add(int(v))
            except Exception:
                continue
        accident_scenario = [t.name for t in accident_types if t.id in norm_ids]

        scene_type = file.scene_type or []
        if "III级响应" in scene_type:
            logger.debug(f"3级响应开始: {chunks_file_path} {accident_scenario}")
            result = run_emergency_3(chunks_file_path, accident_scenario)
        elif "II级响应" in scene_type:
            logger.debug(f"2级响应开始: {chunks_file_path} {accident_scenario}")
            result = run_emergency_32(chunks_file_path, accident_scenario)
        elif "I级响应" in scene_type:
            logger.debug(f"1级响应开始: {chunks_file_path} {accident_scenario}")
            result = run_emergency_321(chunks_file_path, accident_scenario)
        else:
            return JsonModel(code=400, data=None, msg=f"不支持的场景类型").to_response(headers=headers)

        conversation = db_add_conversation(payload.query, dialog_id=dialog.id, references=[], file_id=payload.file_id,
                                           answer=result["final_plan"])
        return JsonModel(code=200,
                         data={**result, "conversation_id": str(conversation.id), "dialog_id": str(dialog.id)},
                         msg=f"应急演练对话成功").to_response(headers=headers)

    except Exception as e:
        logger.error(f"应急演练对话失败: {e}")
        return JsonModel(code=500, data=None, msg=f"应急演练对话失败").to_response(headers=headers)
